/*****************************************************************************
** CS221 Computer Vision Project
** Copyright (c) 2006, Stanford University
**
** FILENAME:    KMeans.h
** AUTHOR(S):  	Chanh Nguyen <chanhn@stanford.edu>
** 				Divij Gupta <divijg@stanford.edu>
** 				Yongwhan Lim <yongwhan@stanford.edu>
** DESCRIPTION:
** This runs the KMeans algorithm over an image to find the K clusters.
**

*****************************************************************************/

#include <cassert>
#include <iostream>
#include <fstream>
#include <algorithm>
#include <fstream>
#include <iostream>

#include "cv.h"
#include "cxcore.h"
#include "highgui.h"

#include "classifier.h"
#include "KMeans.h"

using namespace std;


#define MAX_PIXEL_VALUE 255
#define CENTER_MOVE_THRESHOLD .001 //If the center moves by less than this, it has converged.
#define MAX_ITERATIONS 1000

// Classifier class ---------------------------------------------------------

// default constructor
KMeans::KMeans(const IplImage *sourceimage, int num_clusters_)
{
	image = sourceimage;
	num_clusters = num_clusters_;
	//Find centers
	for(int i = 0; i < 100; i++) {
		float r = rand() % MAX_PIXEL_VALUE;
	}
	for(int i = 0; i < num_clusters; i++) {
		KCluster c;
		float r = rand() % MAX_PIXEL_VALUE;
		c.center[0] = r;
		c.center[1] = r;
		c.center[2] = r;
		r = rand() % MAX_PIXEL_VALUE;
		c.random[0] = r;
		r = rand() % MAX_PIXEL_VALUE;
		c.random[1] = r;
		r = rand() % MAX_PIXEL_VALUE;
		c.random[2] = r;
		clusters.push_back(c);
	}
	converged = false;
}

// destructor
KMeans::~KMeans()
{
}


std::vector<KCluster> KMeans::findClusters(){
	int iters = 0;
    KPoint all_points[image->width][image->height];
	//Read in pixel values
	for(int x = 0; x < image->width; x++) {
		for(int y = 0; y < image->height; y++) {
			uchar b = CV_IMAGE_ELEM(image, uchar, y, x*3 + 0);
			uchar g = CV_IMAGE_ELEM(image, uchar, y, x*3 + 1);
			uchar r = CV_IMAGE_ELEM(image, uchar, y, x*3 + 2);
			KPoint p; p.x = x; p.y = y;
			p.values[0] = b; p.values[1] = g; p.values[2] = r;
			all_points[x][y] = p;
		}
	}

	while(iters < MAX_ITERATIONS && !converged) {
		//cout << "ITER: " << iters << endl;
		for(int i = 0; i < num_clusters; i++) {
			clusters[i].points.clear(); //Clear
		}

		//Assign each pixel to cluster
		for(int x = 0; x < image->width; x++) {
			for(int y = 0; y < image->height; y++) {
				KPoint p = all_points[x][y];
				int cluster_num = assignCluster(p);
				clusters[cluster_num].points.push_back(p);
			}
		}

		converged = true;
		//Find center and clear cluster and check for convergence
		for(int i = 0; i < num_clusters; i++) {
			KCluster *c = &clusters[i];
			for(int color = 0; color < 3; color++) {
				c->prev_center[color] = c->center[color];
			}
			findAverage(c);
			//findFarthest(c);
			if (distance(c->center, c->prev_center) > CENTER_MOVE_THRESHOLD)
				converged = false;
		}

		iters++;
	}
	for(int i = 0; i < num_clusters; i++) {
		KCluster *c = &clusters[i];
		//cout << "ITER: " << iters << "CLUSTER: " << i << "CENTER: " << " DIFF: " << distance(c->center, c->prev_center) << "POINTS: " << c->points.size() << endl;
	}

	//cout << "DONE CLUSTERING" << endl;
	return clusters;

}

void KMeans::findAverage(KCluster *c) {
	for(int color = 0; color < 3; color++) {
		float sum = 0;
		for(int i = 0; i < c->points.size(); i++) {
			sum += c->points[i].values[color];
		}
		if(c->points.size() != 0)
			c->center[color] =  sum/c->points.size();
	}
}

void KMeans::findFarthest(KCluster *c) {
	for(int color = 0; color < 3; color++) {
		float farthest[3] = {0, 0, 0};
		for(int i = 0; i < c->points.size(); i++) {
			float dist = c->points[i].values[color] - c->center[color];
			if(dist > farthest[color]);
			farthest[color] = dist;
		}
		c->center[color] = farthest[color];
	}
}

int KMeans::assignCluster(KPoint &p) {
	float least_distance = distance(p.values, clusters[0].center);
	int nearest_cluster = 0;
	for(int i = 1; i < num_clusters; i++) {
		float dist = distance(p.values, clusters[i].center);
		if(dist < least_distance) {
			least_distance = dist;
			nearest_cluster = i;
		}
	}
	return nearest_cluster;
}

float KMeans::distance(float a[], float b[]) {
	return sqrt(pow((a[0] - b[0]), 2) + pow((a[1] - b[1]), 2) + pow((a[2] - b[2]), 2));
}


IplImage* KMeans::getClusterImage() {
	IplImage* img=cvCreateImage(cvSize(image->width,image->height),IPL_DEPTH_32F,3);
	for(int i = 0; i < num_clusters; i++) {
		KCluster c = clusters[i];
		for(int j = 0; j < c.points.size(); j++) {
			KPoint p = c.points[j];
			/*
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 0]=c.random[0];
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 1]=c.random[1];
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 2]=c.random[2];
			*/
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 0]=c.center[0];
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 1]=c.center[1];
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 2]=c.center[2];

		}
	}
	return img;
}

IplImage* KMeans::getSingleClusterImage(int cluster_num) {
	IplImage* img=cvCreateImage(cvSize(image->width,image->height),IPL_DEPTH_32F,3);
		KCluster c = clusters[cluster_num];
		for(int x = 0; x < image->width; x++) {
			for(int y = 0; y < image->height; y++) {
				((float *)(img->imageData + (y)*img->widthStep))[(x)*img->nChannels + 0]=255;
				((float *)(img->imageData + (y)*img->widthStep))[(x)*img->nChannels + 1]=255;
				((float *)(img->imageData + (y)*img->widthStep))[(x)*img->nChannels + 2]=255;
			}

		}
		for(int j = 0; j < c.points.size(); j++) {
			KPoint p = c.points[j];
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 0]=0;
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 1]=0;
			((float *)(img->imageData + (p.y)*img->widthStep))[(p.x)*img->nChannels + 2]=0;
		}
		return img;
}
